#include "stdafx.h"
#include "X86_InterruptMgr.h"
#include "X86_Tests.h"
#include "Commands.h"
#include "ProcessEnhancer.h"

void dbghelperUnload(IN PDRIVER_OBJECT DriverObject);
NTSTATUS dbghelperCreateClose(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp);
NTSTATUS dbghelperDefaultHandler(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp);
NTSTATUS dbghelperControl(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp);

#ifdef __cplusplus
extern "C" NTSTATUS DriverEntry(IN PDRIVER_OBJECT DriverObject, IN PUNICODE_STRING  RegistryPath);
#endif

X86_Tests* pTest;
char x86_test_stub[sizeof(X86_Tests) + StaticAllocator::StubSize];

NTSTATUS DriverEntry(IN PDRIVER_OBJECT DriverObject, IN PUNICODE_STRING  RegistryPath)
{
	UNICODE_STRING DeviceName,Win32Device;
	PDEVICE_OBJECT DeviceObject = NULL;
	NTSTATUS status;
	unsigned i;
	

	RtlInitUnicodeString(&DeviceName,L"\\Device\\dbghelper0");
	RtlInitUnicodeString(&Win32Device,DBGHELPER_DEVICE_NAME);

	for (i = 0; i <= IRP_MJ_MAXIMUM_FUNCTION; i++)
		DriverObject->MajorFunction[i] = dbghelperDefaultHandler;
	DriverObject->MajorFunction[IRP_MJ_CREATE] = dbghelperCreateClose;
	DriverObject->MajorFunction[IRP_MJ_CLOSE] = dbghelperCreateClose;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = dbghelperControl;
	
    DriverObject->DriverUnload = dbghelperUnload;
    status = IoCreateDevice(DriverObject,
                            0,
                            &DeviceName,
                            FILE_DEVICE_UNKNOWN,
                            0,
                            FALSE,
                            &DeviceObject);
    if (!NT_SUCCESS(status))
        return status;
    if (!DeviceObject)
        return STATUS_UNEXPECTED_IO_ERROR;

    DeviceObject->AlignmentRequirement = FILE_WORD_ALIGNMENT;
    status = IoCreateSymbolicLink(&Win32Device, &DeviceName);

    DeviceObject->Flags &= ~DO_DEVICE_INITIALIZING;

	//InterruptHandle hInterrup = X86_InterruptMgr::Instance()->RegisterInterrupt(0,0,0);
	//pTest = new(&StaticAllocator(x86_test_stub, sizeof(x86_test_stub))) X86_Tests;
	//test.TestIdtTypes();
	//pTest->StartTestHookIdt0();
	
	
	return status;
}

void dbghelperUnload(IN PDRIVER_OBJECT DriverObject)
{
	UNICODE_STRING Win32Device;
	RtlInitUnicodeString(&Win32Device,L"\\DosDevices\\dbghelper0");
	IoDeleteSymbolicLink(&Win32Device);
	IoDeleteDevice(DriverObject->DeviceObject);
	ProcessEnhancer::Unload();
	//pTest->StopTestHookIdt0();
	//delete pTest;
}


NTSTATUS dbghelperDefaultHandler(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
	Irp->IoStatus.Status = STATUS_NOT_SUPPORTED;
	Irp->IoStatus.Information = 0;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return Irp->IoStatus.Status;
}
NTSTATUS dbghelperCreateClose(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
	Irp->IoStatus.Status = STATUS_SUCCESS;
	Irp->IoStatus.Information = 0;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return Irp->IoStatus.Status;
}

NTSTATUS dbghelperControl(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp)
{
	NTSTATUS status;
	PIO_STACK_LOCATION pStackLocation;
	PDbgProcessEnhance pProcessEnhance;

	pStackLocation = IoGetCurrentIrpStackLocation(Irp);
	switch(pStackLocation->Parameters.DeviceIoControl.IoControlCode)
	{
	case DBGHELPER_IOCTL_PROCESS_ENHANCE:
		if(pStackLocation->Parameters.DeviceIoControl.InputBufferLength != sizeof(DbgProcessEnhance))
		{
			KdPrint(("DbgHelper!dbghelperControl==>Invalid Input Len for process enhance!,expect %d but got %d\n",
				sizeof(DbgProcessEnhance), pStackLocation->Parameters.DeviceIoControl.InputBufferLength));
			status = STATUS_INVALID_PARAMETER;
			break;
		}
		/*
		* TODO
		*/
		pProcessEnhance = (PDbgProcessEnhance)Irp->AssociatedIrp.SystemBuffer;
		if(!ProcessEnhancer::Instance()->EnhanceProcess(pProcessEnhance->PId))
		{
			status = STATUS_UNSUCCESSFUL;
			break;
		}
		break;
	default:
		KdPrint(("DbgHelper!dbghelperControl==>Unrecognised control code 0x%08x\n", pStackLocation->Parameters.DeviceIoControl.IoControlCode));
		status = STATUS_INVALID_PARAMETER;
		break;
	}

	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = 0;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return Irp->IoStatus.Status;
}
